from multiprocessing import freeze_support

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, models, transforms
from tqdm import tqdm


def main():
    # 数据预处理
    transform = transforms.Compose(
        [
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )

    # 加载数据集
    train_dataset = datasets.ImageFolder(root="data/TrainingSet", transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)

    # 获取类别数量
    num_classes = len(train_dataset.classes)

    # 加载预训练的ResNet50模型
    model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
    model.fc = nn.Linear(2048, num_classes)  # 修改最后的全连接层

    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 训练参数
    num_epochs = 5
    best_acc = 0.0
    best_model_path = "best_resnet50_model.pth"
    train_losses = []
    val_accuracies = []

    # 训练循环
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        epoch_loss = 0.0

        # 使用tqdm显示进度条
        with tqdm(train_loader, unit="batch") as tepoch:
            for inputs, labels in tepoch:
                tepoch.set_description(f"Epoch {epoch + 1}/{num_epochs}")

                inputs = inputs.to(device)
                labels = labels.to(device)

                # 前向传播
                outputs = model(inputs)
                loss = criterion(outputs, labels)

                # 反向传播和优化
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                running_loss += loss.item()
                epoch_loss += loss.item()

                # 更新进度条
                tepoch.set_postfix(loss=running_loss / (tepoch.n + 1))

        # 记录训练损失
        train_losses.append(epoch_loss / len(train_loader))

        # 验证阶段
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in train_loader:  # 使用训练集进行验证
                inputs = inputs.to(device)
                labels = labels.to(device)
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        val_acc = correct / total
        val_accuracies.append(val_acc)
        print(f"Epoch [{epoch + 1}/{num_epochs}], Validation Accuracy: {val_acc:.4f}")

        # 保存最佳模型
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), best_model_path)
            print(f"New best model saved with accuracy: {best_acc:.4f}")

    # 绘制训练曲线
    plt.figure(figsize=(12, 5))

    # 绘制训练损失曲线
    plt.subplot(1, 2, 1)
    plt.plot(range(1, num_epochs + 1), train_losses, label="Train Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title("Training Loss")
    plt.legend()

    # 绘制验证准确率曲线
    plt.subplot(1, 2, 2)
    plt.plot(range(1, num_epochs + 1), val_accuracies, label="Validation Accuracy")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy")
    plt.title("Validation Accuracy")
    plt.legend()

    # 保存并显示图像
    plt.tight_layout()
    plt.savefig("training_curve.png")
    plt.show()

    print(f"训练完成，最佳模型已保存为{best_model_path}，准确率: {best_acc:.4f}")


if __name__ == "__main__":
    freeze_support()
    main()
